import matplotlib.pyplot as plt
import pandas as pd
from statsmodels.api import tsa
import tushare as ts
import datetime as dt
import warnings
warnings.filterwarnings('ignore', 'statsmodels.tsa.arima_model.ARMA', FutureWarning)
warnings.filterwarnings('ignore', 'statsmodels.tsa.arima_model.ARIMA', FutureWarning)

plt.rc("font",family='MicroSoft YaHei',weight="bold")


class ArmaTrainer():
    def initData(self,
                 code='000001.SZ',
                 stockType='daily',
                 start=str(dt.datetime.today() - dt.timedelta(days=200)),
                 end=str(dt.datetime.today() + dt.timedelta(days=1))):
        # data = ts.get_hist_data(code)
        pro = ts.pro_api()
        data = pro.query('daily', ts_code=code)
        # data = pro.query('daily', ts_code='000001.SZ', start_date='20180701', end_date='20180718')
        # data.index = pd.to_datetime(data.index)
        # data = data.sort_index()
        return data

    def predict(self, data, param):
        # 预测模型
        arma = tsa.ARMA(data[[param]], (1,1) )
        # 训练
        model = arma.fit()
        # 预测
        predict = model.predict(len(data))
        return predict[len(data)]

    def predictNext(self, data):
        open = self.predict(data, 'open')
        close = self.predict(data, 'close')
        high = self.predict(data, 'high')
        low = self.predict(data, 'low')
        vol = self.predict(data, 'vol')
        nextDate = data['trade_date'][0]
        formatNextDate = dt.datetime.strptime(nextDate, '%Y%m%d') + dt.timedelta(1)
        data = data.append([{
            'ts_code': data['ts_code'][0],
            'trade_date': formatNextDate.strftime('%Y%m%d'),
            'vol': vol,
            'open': open,
            'close': close,
            'high': high,
            'low': low
        }], ignore_index=True)
        data.rename(columns={'trade_date': 'datetime', 'vol': 'volume'}, inplace=True)
        # data.set_index(['datetime'], inplace=True)
        # data.index = pd.to_datetime(data.index)
        # data = data.sort_index()
        data.sort_values("datetime", inplace=True, ignore_index=True)  # 排序
        return data


if __name__ == '__main__':
    at = ArmaTrainer()
    data = at.initData('000001.SZ')
    print(data)
    nextData = at.predictNext(data)
    print(nextData)